import os
from dotenv import load_dotenv, find_dotenv

from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.prompts import PromptTemplate
from langchain.agents import AgentExecutor, create_self_ask_with_search_agent

_ = load_dotenv(find_dotenv())

# from langchain_community.chat_models import ChatZhipuAI 
# model = ChatZhipuAI(
#                   model=model,
#                   temperature=0.9,              
#                 )

from langchain_openai import ChatOpenAI
model = ChatOpenAI(temperature=0)

tavily_api_key = os.environ.get('TAVILY_API_KEY')
if tavily_api_key is None:
    raise ValueError("tavily_api_key is not set in the .env file")

search = TavilySearchResults(max_results=1, name="Intermediate Answer")

tools = [search]
 


template = '''
Question: Who lived longer, Muhammad Ali or Alan Turing?
Are follow up questions needed here: Yes.
Follow up: How old was Muhammad Ali when he died?
Intermediate answer: Muhammad Ali was 74 years old when he died.
Follow up: How old was Alan Turing when he died?
Intermediate answer: Alan Turing was 41 years old when he died.
So the final answer is: Muhammad Ali

Question: When was the founder of craigslist born?
Are follow up questions needed here: Yes.
Follow up: Who was the founder of craigslist?
Intermediate answer: Craigslist was founded by Craig Newmark.
Follow up: When was Craig Newmark born?
Intermediate answer: Craig Newmark was born on December 6, 1952.
So the final answer is: December 6, 1952

Question: Who was the maternal grandfather of George Washington?
Are follow up questions needed here: Yes.
Follow up: Who was the mother of George Washington?
Intermediate answer: The mother of George Washington was Mary Ball Washington.
Follow up: Who was the father of Mary Ball Washington?
Intermediate answer: The father of Mary Ball Washington was Joseph Ball.
So the final answer is: Joseph Ball

Question: Are both the directors of Jaws and Casino Royale from the same country?
Are follow up questions needed here: Yes.
Follow up: Who is the director of Jaws?
Intermediate answer: The director of Jaws is Steven Spielberg.
Follow up: Where is Steven Spielberg from?
Intermediate answer: The United States.
Follow up: Who is the director of Casino Royale?
Intermediate answer: The director of Casino Royale is Martin Campbell.
Follow up: Where is Martin Campbell from?
Intermediate answer: New Zealand.
So the final answer is: No

Question: {input}
Are followup questions needed here:{agent_scratchpad}
'''

#prompt = hub.pull("hwchase17/self-ask-with-search")
prompt = PromptTemplate.from_template(template).partial(tools=tools)

agent = create_self_ask_with_search_agent(model, tools, prompt)

agent_executor = AgentExecutor(
    agent=agent, tools=tools, verbose=True, handle_parsing_errors=True
)

agent_executor.invoke({"input": "《雷雨》作者的家乡是哪里?"})
#agent_executor.invoke({"input": "What is the hometown of the reigning men's U.S. Open champion?"})